package com.rcy.server.wrapper;

import org.springframework.web.bind.annotation.RequestMethod;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.util.Arrays;
import java.util.stream.Collectors;

/**
 * @author RuicyWu
 * 用来重新包装HttpSevletRequest对象
 * 声明一个带有HttpServletRequest入参的构造器，
 * 从该参数对象的流中解析数据，如果没有则继续从parameterMap中获取，
 * 然后以key=value&key=value形式拼接。用streamBody接收。
 * 重写getInputStream方法，以后每次调用getInputStream方法，
 * 其实是重新利用streamBody重新new一个流，所以可以多次读取。
 */
public class InputStreamHttpServletRequestWrapper extends HttpServletRequestWrapper {
 
 
    private final byte[] streamBody;
    private static final int BUFFER_SIZE = 2048;
 
   
    public InputStreamHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        byte[] bytes = inputStream2Byte(request.getInputStream());
        if (bytes.length == 0 && RequestMethod.POST.name().equals(request.getMethod())) {
            //从ParameterMap获取参数，并保存以便多次获取
            bytes = request.getParameterMap().entrySet().stream()
                    .map(entry -> {
                        String result;
                        String[] value = entry.getValue();
                        if (value != null && value.length > 1) {
                            result = Arrays.stream(value).map(s -> entry.getKey() + "=" + s)
                                    .collect(Collectors.joining("&"));
                        } else {
                            result = entry.getKey() + "=" + value[0];
                        }
 
                        return result;
                    }).collect(Collectors.joining("&")).getBytes();
        }
 
        streamBody = bytes;
    }
 
    private byte[] inputStream2Byte(InputStream inputStream) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        byte[] bytes = new byte[BUFFER_SIZE];
        int length;
        while ((length = inputStream.read(bytes, 0, BUFFER_SIZE)) != -1) {
            outputStream.write(bytes, 0, length);
        }
 
        return outputStream.toByteArray();
    }


    /**
     * 重写getInputStream()方法,使其每次都新建一个流
     * @return 返回新的流
     * @throws IOException
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        ByteArrayInputStream inputStream = new ByteArrayInputStream(streamBody);
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }
 
            @Override
            public boolean isReady() {
                return false;
            }
 
            @Override
            public void setReadListener(ReadListener listener) {
 
            }
 
            @Override
            public int read() throws IOException {
                return inputStream.read();
            }
        };
    }

    /**
     * 重写获取字符流的方法,效果同上
     * @return
     * @throws IOException
     */
    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }
}